#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2024/11/21 16:34
# @Author : XXX
# @Site : 
# @File : model_class.py
# @Software: PyCharm

# 定义模型
import torch
from torch import nn


class Model(nn.Module):
    def __init__(self, len_state=6):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(len_state, 50)
        self.linear2 = nn.Linear(50, 50)
        self.linear3 = nn.Linear(50, 50)
        self.linear4 = nn.Linear(50, 50)
        self.linear5 = nn.Linear(50, 50)
        self.linear6 = nn.Linear(50, 1)

    def forward(self, x):
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        x = torch.relu(self.linear3(x))
        x = torch.relu(self.linear4(x))
        x = torch.relu(self.linear5(x))
        x = self.linear6(x) # 最后一层不需要激活函数
        return x